# -*- coding: utf-8 -*-

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

def read_data_from_path(path):
    '''
    根据给出的 path，读取 .xyz 文件
    '''
    with open(path,'r') as file_object:
        # 将文件一列一列读取，并放到一个叫 contents 的变量中
        contents = file_object.readlines()
        # num 为文件中的行数。
        num = len(contents)
        
        # 构建一个临时变量，用来存取 Au20 的20原子的数据
        cluster = []
        # 第一行是 Au20
        cluster.append(float(contents[0].rstrip()))
        # 第二行是 能量
        cluster.append(float(contents[1].rstrip()))
        
        # 读取剩余的行，即 20 个AU20 原子的坐标
        for i in range(2,num):
            # 将每一行中多余的空格剔除
            str_space_limit = ' '.join(contents[i].split())
            # 将每行以空格为分隔符拆分，并保存到一个叫 au 的临时变量中
            au = [float(i) for i in str_space_limit.split(' ')[1:]]
            # 将 au 添加到 cluster 中
            cluster += au
        
        au_num = cluster[0]
        # 读取能量数据
        energy = cluster[1]
        # 将坐标数据作为一个向量输出
        cdn_vec = np.array(cluster[2:])
        # 将坐标数据作为一个
        cdn_mat = cdn_vec.reshape((-1,3)) #Matrix
        
    return au_num, energy, cdn_vec, cdn_mat

def plot_cluster(coordination):
    '''
    画出 AU20 团簇的 3D 结构图
    '''
    fig = plt.figure(figsize=(4,4))
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(xs = coordination[:,0], ys=coordination[:,1], zs=coordination[:,2])
    plt.show()


def plot_surface(coordination):
    '''
    画出 AU20 团簇的 3D surface 图。
    不过貌似不太直观和准确。
    '''
    fig = plt.figure()
    ax = plt.figure().add_subplot(projection='3d')
    ax.plot_trisurf(coordination[:,0], coordination[:,1], coordination[:,2]\
                    , linewidth=0.2, antialiased=True)
    plt.show()

if __name__ == '__main__':
    # 测试，可以不用运行
    path = r'..\题目\附件\Au20_OPT_1000\0.xyz'
    au_num, energy, cdn_vec, cdn_mat = read_data_from_path(path)
    print(f'原子个数为:{au_num}')
    print(f'能量为：{energy}')
    print(f'坐标向量：{cdn_vec}')
    print(f'坐标矩阵:{cdn_mat}')
    plot_cluster(cdn_mat)
    plot_surface(cdn_mat)
    
    
    
